import random
from deap import base, creator, tools, algorithms
import statistics
from geometry_process import *
import os
import seaborn as sns
from matplotlib.colors import Normalize, ListedColormap
import copy
# 使用Seaborn的color_palette功能生成50种颜色
colors = sns.color_palette("hsv", 50)

# 创建自定义的颜色映射
custom_cmap = ListedColormap(colors)
global res_dict



import matplotlib.pyplot as plt
from matplotlib.colors import Normalize, ListedColormap
import matplotlib.cm as cm
from matplotlib.patches import Rectangle

from collections import Counter
import seaborn as sns

# 使用Seaborn的color_palette功能生成50种颜色
colors = sns.color_palette("hsv", 50)

# 创建自定义的颜色映射
custom_cmap = ListedColormap(colors)

# 绘画一个Multipolygon, ind用于指示其所属的模块
def draw_individual_multipolygon(individual):
    plt.axis('equal')
    cmap = plt.cm.Reds
    norm = plt.Normalize(vmin=0, vmax=len(individual))
    for i, ind in enumerate(individual):
        if len(ind.block) == 0:
            continue
        polygon = ind.get_geoms()
        x, y = polygon.exterior.xy

        plt.plot(x, y, 'o', color='#f16824', zorder=1)
        plt.fill(x, y, color=cmap(norm(i)), alpha=0.5, zorder=2)
        
# 绘画一个Multipolygon, ind用于指示其所属的模块type
def draw_type_multipolygon(individual):
    plt.axis('equal')
    cmap = cm.get_cmap('tab20')  # 使用有更多颜色的colormap
    dimensions = []
    dimension_colors = {}

    # 为每个独特的长宽组合分配颜色
    for ind in individual:
        if len(ind.block) == 0:
            continue

        width, length = get_width_length(ind.get_geoms())
        dimension = (round(width, 2), round(length, 2))
        if dimension not in dimensions:
            dimensions.append(dimension)

    norm = Normalize(vmin=0, vmax=len(dimensions) - 1)
    for i, dimension in enumerate(dimensions):
        dimension_colors[dimension] = cmap(norm(i))

    # 绘制每个模块
    for ind in individual:
        if len(ind.block) == 0:
            continue

        polygon = ind.get_geoms()
        width, length = get_width_length(polygon)
        dimension = (round(width, 2), round(length, 2))
        color = dimension_colors[dimension]
        x, y = polygon.exterior.xy
        plt.fill(x, y, color=color, alpha=0.5, zorder=2)

    plt.show()


def draw_modules_by_type(individual):
    # 首先统计每种类型的模块数量
    type_count = Counter()
    for ind in individual:
        if len(ind.block) == 0:
            continue
        width, length = get_width_length(ind.get_geoms())
        type_count[(round(width, 2), round(length, 2))] += 1
    
    # 按类型数量排序，以便整齐排列
    sorted_types = sorted(type_count.items(), key=lambda x: x[1], reverse=True)

    # 创建绘图区域
    fig, ax = plt.subplots(figsize=(12, 8))
    ax.set_xlim(0, 100)
    ax.set_ylim(0, 100)
    
    # 定义初始位置
    current_x = 0
    current_y = 0
    max_height_in_row = 0
    
    # 绘制每种模块类型
    for (width, length), count in sorted_types:
        # 如果当前行放不下，换到下一行
        if current_x + width > ax.get_xlim()[1]:
            current_x = 0
            current_y += max_height_in_row + 10
            max_height_in_row = 0
        
        # 绘制模块
        ax.add_patch(Rectangle((current_x, current_y), width, length, edgecolor='black', facecolor='skyblue', alpha=0.5))
        # 在模块旁边添加数量标签
        ax.text(current_x + width + 1, current_y + length / 2, f"×{count}", verticalalignment='center')
        # 在模块旁边添加尺寸标签
        ax.text(current_x + width + 1, current_y, f"{width}x{length}", verticalalignment='center')
        
        # 更新位置
        current_x += width + 15
        max_height_in_row = max(max_height_in_row, length)

    plt.axis('off')  # 关闭坐标轴
    plt.show()

# 生成实例的代码，返回一个列表module_list
def my_indices(res_dict):
    print(f"Creating individual in process {os.getpid()}")
    module_list = []
    # 生成与初始blocks数量相当的modules
    module_list = [Module() for _ in range(len(res_dict))]
    # 打乱原始的block list，使得初始种群尽可能多样化
    shuffled_block_list = np.random.permutation(len(res_dict))
    for each_block in shuffled_block_list:
        for each_module in module_list:
            # 将每一个block塞入第一个可行的module中
            each_module.add_block([each_block])
            if each_module.is_feasible():
                break
            else:
                each_module.delete_block([each_block])
    return module_list

# 估计实例fitness的函数
def eval_func(individual):
    module_num = 0
    diff_module_num = 0
    dimensions = []

    for module in individual:
        if not module.is_feasible():
            return (1e9, 1e9,)
        if len(module.block) != 0:
            module_num += 1

            # 检查同种模块数量
            # 获取模块的长和宽，并确保长和宽是排序过的
            dimension = get_width_length(module.get_geoms())
            if dimension not in dimensions:
                diff_module_num += 1
                dimensions.append(dimension)
            
    return (module_num, diff_module_num,)

# 两点交叉函数，用于交配生成子代
def my_cxTwoPoint(ind1, ind2):
    size = min(len(ind1), len(ind2))
    cxpoint1 = random.randint(1, size)
    cxpoint2 = random.randint(1, size - 1)
    if cxpoint2 >= cxpoint1:
        cxpoint2 += 1
    else:  # Swap the two cx points
        cxpoint1, cxpoint2 = cxpoint2, cxpoint1
    
    ind1[cxpoint1:cxpoint2], ind2[cxpoint1:cxpoint2] = ind2[cxpoint1:cxpoint2], ind1[cxpoint1:cxpoint2]
    
    ind1 = refine_gene(ind1, ind1[cxpoint1:cxpoint2])
    ind2 = refine_gene(ind2, ind2[cxpoint1:cxpoint2])
    return ind1, ind2,

# 整理由于两点交换基因导致的基因重复、基因缺失情况
def refine_gene(ind, par_ind):
    blocks = set()
    par_blocks = set()

    for each in par_ind:
        par_blocks.update(each.block)
    
    new_ind = []

    for each in ind:
        if each in par_ind:
            new_ind.append(each)
            continue
        ind_blocks = set(each.block)
        # 如果和伴侣基因片段包含的block有重复，则删除此基因，保证能够保留来自伴侣的基因特征
        if ind_blocks & par_blocks:
            continue
        else:
            new_ind.append(each)

        
    while len(new_ind) < len(ind):
        new_ind.append(Module())

    # 获得已有的所有blocks，已经缺失的blocks
    for each in new_ind:
        blocks.update(each.block)

    miss_blocks = set(res_dict.keys()) - blocks

    # 将缺失的blocks插入第一个可行的Module
    for each_block in miss_blocks:
        for each_module in new_ind:

            each_module.add_block([each_block])
            if each_module.is_feasible():

                break
            else:
                each_module.delete_block([each_block])
    return creator.Individual(new_ind)

# 变异函数
def my_mut(ind, indpb):

    new_ind = []
    blocks = set()
    temp_ind = []
    # 随机删除一些modules
    for i in range(len(ind)):
        if random.random() >= indpb:
            new_ind.append(ind[i])
            temp_ind.append(copy.deepcopy(ind[i]))
    # 将空modules重新加入list以保证长度相等
    while len(new_ind) < len(ind):
        new_ind.append(Module())
        temp_ind.append(Module())
    # 获得已有的所有blocks，已经缺失的blocks
    for each in new_ind:
        blocks.update(each.block)

    miss_blocks = set(res_dict.keys()) - blocks

    for each_block in miss_blocks:
        for each_module in new_ind:

            each_module.add_block([each_block])
            if each_module.is_feasible():

                break
            else:
                each_module.delete_block([each_block])
    return creator.Individual(new_ind),